import subprocess

def get_gpu_status():
    # 调用 nvidia-smi 命令
    result = subprocess.run(
        ["nvidia-smi", "--query-gpu=index,name,utilization.gpu,memory.total,memory.used,temperature.gpu", "--format=csv,noheader,nounits"],
        stdout=subprocess.PIPE,
        text=True
    )

    # 解析输出
    gpu_data = []
    for line in result.stdout.strip().split('\n'):
        index, name, gpu_util, mem_total, mem_used, temp = line.split(', ')
        gpu_data.append({
            "index": int(index),
            "name": name,
            "gpu_utilization": int(gpu_util),
            "memory_total": int(mem_total),
            "memory_used": int(mem_used),
            "temperature": int(temp)
        })

    return gpu_data

# 打印 GPU 状态
gpu_status = get_gpu_status()
for gpu in gpu_status:
    print(f"GPU {gpu['index']} ({gpu['name']}):")
    print(f"  GPU Utilization: {gpu['gpu_utilization']}%")
    print(f"  Memory Used: {gpu['memory_used']} MiB / {gpu['memory_total']} MiB")
    print(f"  Temperature: {gpu['temperature']}°C")
    print()